We could drop the select to simplify this a bit

We could also mostly support the context cancellation by wrapping the in io.Reader we pass into DisplayJSONMessagesToStream with something like:

type stopableReader struct {
    r io.Reader
    closed error
}

func (c *stopableReader) Read(p []byte) (int, error) {
  if (c.closed != nil) {
    return 0, c.closed
  }
  return c.r.Read(p)
}

and then set closed = ctx.Err() after the context is done.

I think i'd prefer we be explicit about it instead of in the io.Reader. I'll remove the select

The io.Reader wrapper is a separate thing though.

With the current implementation the context cancellation isn't actually handled once the DisplayJSONMessagesToStream is called, the cancellation won't have any actual effect.

I pushed a small test to your branch to illustrate that.

The reader wrapper is trying to achieve that, that means this would effectively remove the need to make changes to the pkg/jsonmessage on moby side.

(sorry, late to comment; I saw the discussion, and was indeed wondering if closing the reader would make the jsonmessages properly cancel, as ISTR it was originally designed around some of that)

With the current implementation the context cancellation isn't actually handled once the DisplayJSONMessagesToStream is called, the cancellation won't have any actual effect.

yes, but the function was never designed for that. This PR is only trying to solve returning a context error instead of just a nil (when the reader is implicitly cancelled by context through say http.NewRequestWithContext).

The reader wrapper is trying to achieve that, that means this would effectively remove the need to make changes to the pkg/jsonmessage on moby side.

This PR's intent is to circumvent the moby side PR - or at least find an alternative without needing to refactor the function.

and was indeed wondering if closing the reader would make the jsonmessages properly cancel

Yes you are correct, it won't. The jsonmessages functions rely on the fact that io.EOF was returned to break the for loop.

EDIT: I meant to say that jsonmessages functions solely rely on io.EOF from the reader to exit, but that doesn't mean we will know why it was cancelled (the io.EOF reason). So closing the reader would exit the loop. https://github.com/moby/moby/blob/master/pkg/jsonmessage/jsonmessage.go#L237-L242

It will also exit on any Read other error and will return it:

https://github.com/moby/moby/blob/1153242d3a6c3635770db65d9233bb892ee1cbc3/pkg/jsonmessage/jsonmessage.go#L237-L243

If the Read inside Decode returns context.Canceled then it would also be returned from Decode.
That's what the stopableReader above does.

This PR's intent is to circumvent the moby side PR - or at least find an alternative without needing to refactor the function.

I get that, but why not make the DisplayStream function actually respect the cancellation if we can?
We control the reader passed so we can make it start returning error for subsequent Reads after the context is done.

With that stopableReader I posted above, we could just wait for ctx.Done() channel to be closed and set r.closed = ctx.Err() to make the reader start returning the context error.

i see, alright I'll look into this then

As I mentioned in #5645 (comment) and #5645 (comment), I don't think adding a context to jsonmessage functions is the best approach – it would be simpler to let those continue only caring about whether the read fails or not to know when to stop instead of having to care about both things.

@vvoland's suggestion allows addressing the issue without needless changes to jsonmessage.

One other thing consider is to add 2 separate functions; one that takes an auxCallBack and one without; looking at the CLI code, I see that we need the auxCallBack in only 2 places;

cli/command/container/create.go:        return jsonmessage.DisplayJSONMessagesToStream(responseBody, out, nil)
cli/command/image/import.go:    return jsonmessage.DisplayJSONMessagesToStream(responseBody, dockerCli.Out(), nil)
cli/command/image/load.go:              return jsonmessage.DisplayJSONMessagesToStream(response.Body, dockerCli.Out(), nil)
cli/command/image/push.go:              err = jsonmessage.DisplayJSONMessagesToStream(responseBody, streams.NewOut(io.Discard), handleAux(dockerCli))
cli/command/image/push.go:      return jsonmessage.DisplayJSONMessagesToStream(responseBody, dockerCli.Out(), handleAux(dockerCli))
cli/command/image/trust.go:             if err := jsonmessage.DisplayJSONMessagesToStream(in, ioStreams.Out(), nil); err != nil {
cli/command/image/trust.go:     if err := jsonmessage.DisplayJSONMessagesToStream(in, ioStreams.Out(), handleTarget); err != nil {
cli/command/image/trust.go:     return jsonmessage.DisplayJSONMessagesToStream(responseBody, out, nil)
cli/command/plugin/install.go:  if err := jsonmessage.DisplayJSONMessagesToStream(responseBody, dockerCli.Out(), nil); err != nil {
cli/command/plugin/push.go:     return jsonmessage.DisplayJSONMessagesToStream(responseBody, dockerCli.Out(), nil)
cli/command/plugin/upgrade.go:  if err := jsonmessage.DisplayJSONMessagesToStream(responseBody, dockerCli.Out(), nil); err != nil {
cli/command/service/helpers.go: err := jsonmessage.DisplayJSONMessagesToStream(pipeReader, dockerCli.Out(), nil)
cli/command/swarm/ca.go:        err := jsonmessage.DisplayJSONMessagesToStream(pipeReader, dockerCli.Out(), nil)

Alternatively, we could consider (for a expandable UX) having it accept a DisplayOptions{} struct (or fancy functional DisplayOptions()) , e.g.

err := DisplayWithOptions(ctx, in, stream, Options{
    AuxCallback: myCallback,
})

another is to make Display a variadic function like so:

type Options func(*options)

type options struct {
	AuxCallback func(JSONMessage)
}

func WithAuxCallback(cb func(JSONMessage)) Options {
	return func(o *options) {
		o.AuxCallback = cb
	}
}

func Display(ctx context.Context, in io.Reader, stream Stream, opts ...Options) error {
}

Then when you use it, you don't even need to specify the last parameter in most cases.

err := jsonstream.Display(ctx, reader, out)
// or
err := jsonstream.Display(ctx, reader, out, jsonstream.WithAuxCallback(func(j jsonstream.JSONMessage){}))

If we could name this package jsonstream it could help disambiguate with the jsonmessage package, and we could name the function Display() as it would be called as jsonstream.Display()

This is still an internal package, so not too problematic, but it looks like JSONError is literally used in a single place in the CLI. We may want to have a look if we want to keep it as an exported type, or if there's alternatives to reduce the "public" (but right now, still "internal") API;

The idea behind this was to keep all the jsonmessage code in a singular place inside an internal package inside the CLI. Maybe i'm being a bit excessive, but it felt cleaner